[ROCm] Add FlyDSL paged-attention decode backend (dense + fp8) with benchmark by avbokovoy · Pull Request #463 · meta-pytorch/MSLK

avbokovoy · 2026-07-29T10:43:02Z

Summary

Replaces the CK paged-attention decode operators with a FlyDSL backend and adds a native-fp8 decode path, wired into the existing fmha op registry (flydsl_decoder / flydsl_splitk) with backward-compatible op aliases. FlyDSL becomes the sole decode backend across gfx942 and gfx950; the CK operator path is removed.

Kernels

Dense f16/bf16 decode uses three kernels selected by shape: pa_decode_gfx950 (head-packed MFMA + double-buffered wide V load, ds_read_tr16_b64; gfx950, GQA ratio 1..16), pa_decode_gfx950_coop (gfx950 cooperative-DMA for ratios that cannot head-pack), and pa_decode_generic (arch-generic per-warp fallback for gfx942 and off-gfx950). pa_decode_fp8 adds native fp8 (e4m3fn + symmetric per-token scale) paged decode covering MQA/GQA, D=128/256, and any context length; it is opt-in per call via Inputs.quantize_kv_to_fp8 and is CUDA-graph capturable. All kernels (dense, split-K reduce, and fp8) are AOT-registered and precompiled into the bundled cache.

Performance

Measured on gfx950 (MI350X). Dense f16/bf16: FlyDSL is 1.6-10x faster than the old CK decoder (CK uses no matrix cores) and within 1.05-1.37x of Triton (near parity; Triton is marginally faster on raw kernel time). Native fp8: FlyDSL is 1.8-2.5x faster than Triton fp8 and also beats Triton dense on most shapes. Correctness is verified against a torch reference for all backends.

Other changes

Fixes the Triton fp8 decode to use OCP e4m3fn on gfx950 (was hardcoded fnuz, which mis-decoded gfx950 caches). Aligns FlyDSL helper imports with the mslk.flydsl package layout. Adds bench/attn/decoder_bench.py (FlyDSL vs Triton, dense and fp8) with eager timing via the shared do_bench, real CUDA-graph timing with empty-graph detection, and per-shape subprocess isolation. Adds ROCm CI path triggers for the decode backend and its tests.

Test plan

test_flydsl_fp8_decoder (108 passed on gfx950)
Dense f16/bf16 decode correctness vs torch reference
AOT precompile of all dense and fp8 configs on gfx950

meta-codesync · 2026-08-05T19:05:08Z

@q10 has imported this pull request. If you are a Meta employee, you can view this in D114929079.

cthi

Thanks, I left some high level questions. Mainly wondering about the comments arounds the benchmarking.

cthi · 2026-08-07T13:31:44Z

+  * eager (``--no-cuda-graph``, default): shared ``mslk.bench.common.utils.do_bench``.
+  * graph (``--cuda-graph``): HIP graph capture + replay (removes launch overhead).
+
+gfx950 gotchas (see the runners / _bench_ms_graph for detail):


Is these "gotchas" for graph benchmark specific stuff? Or more like limitations we should follow up on? Graph is actually more important than eager, so we should always default our benchmarking/optimization/enablement to be graph-enabled first.

It should be fixed in the latest commit. graph capture now works correct. Adjusted the comments accordingly

cthi · 2026-08-07T13:46:07Z

+
+"""Shared low-level FlyDSL helpers for attention kernels.
+
+Only pip `flydsl==0.2.2` is imported (no ~/FlyDSL/kernels imports).


I don't understand this comment, is it left over from a debugging?

Yes, leftover. Removed that and refactored/trimmed all the comments in this PR

cthi · 2026-08-07T13:46:47Z

+
+# pyre-strict
+
+"""Shared low-level FlyDSL helpers for attention kernels.


These don't look attention specific, should they be moved into some more common flydsl/utils.py?

I did some refactoring in the latest commit, moving it into kernel_intrinsics.py inside flydsl submodule. utils.py now just forwards those primitives

cthi · 2026-08-07T13:49:59Z

+  * fp8 runners cache the ``flyc.compile`` CompiledFunction so timing is kernel-only
+    and scales with KV (calling the public dispatcher directly pays ~0.38ms/call of
+    JIT dispatch that hides the ~0.02ms kernel — flat, meaningless numbers).
+  * flydsl_fp8 IS graph-capturable (its kernels thread the capture stream); dense


Can you explain this "empty graph" issue? It seems we modify a lot of the benchmarking to handle this.

It should be fixed in the latest commit. graph capture now works correct. Adjusted the comments accordingly

Replace the CK paged-attention decode operators with FlyDSL kernels and add a native-fp8 decode path, wired into the fmha op registry (flydsl_decoder / flydsl_splitk). Kernels (mslk/attention/fmha/flydsl/): - pa_decode_gfx950: primary head-packed MFMA fast path (double-buffered wide V load, ds_read_tr16_b64), gfx950, GQA ratio in [1,16]. - pa_decode_gfx950_coop: gfx950 cooperative-DMA kernel for ratios the primary path can't head-pack. - pa_decode_generic: arch-generic per-warp fallback (gfx942 + off-gfx950). - pa_decode_fp8: native-fp8 (e4m3fn + symmetric per-token scale) paged decode, with a per-call quantizing adapter and a guarded public dispatcher. MQA/GQA, D=128/256, any context length. Stream threaded through compute + reduce launches so the kernel is CUDA-graph capturable. fp8-KV is opt-in per call via Inputs.quantize_kv_to_fp8. - pa_decode_reduce: split-K partial combine. - utils: shared low-level FlyDSL helpers (WARP_SIZE, dpp/wave-reduce, exp2/rcp/max). AOT: dense (generic/gfx950/coop), split-K reduce, and native-fp8 kernels are all registered in mslk/flydsl/aot.py and precompiled into the bundled cache. Also: - Fix the Triton fp8 decode to use OCP e4m3fn on gfx950 (not fnuz). - Import FlyDSL helpers from the mslk.flydsl package (common/jit), matching the upstream flash-attention layout. - bench/attn/decoder_bench.py: FlyDSL + Triton dense/fp8 backends, eager timing via the shared do_bench, real CUDA-graph timing with empty-graph detection, and subprocess isolation per shape. - ROCm CI path triggers for the decode backend + tests.

avbokovoy · 2026-08-10T12:31:09Z

@cthi Thanks for all the comments on this PR. Let me know if you have any other questions

cthi

Given the PR is quite large, I would suggest we split it into 3 different PRs to get it in:

Changes to existing triton_splitK for gfx950
Addition of new FlyDSL kernel
Integration of new FlyDSL kernel into fMHA as a separate backend (could be merged with above PR). But no removal of CK backend yet.
Removal of CK backend (at a future date)

The reason is that enabling and verifying the new FlyDSL backend for fmha may need some changes internally in Meta, so adding + removing CK in a single PR will make it difficult to merge internally. Thanks.

cc @bottler as well

cthi · 2026-08-11T13:24:28Z

    use_fp32_scales: bool = False
    num_splits: int = 0
+    # Opt-in for FlyDSL native-fp8 paged decode: quantize dense KV to fp8 (e4m3fn)
+    # per call and run the fp8 kernel. gfx950 + G=1 only, else falls back to dense.


Im not sure I understand this feature, why would someone want to quantize the KV on the fly per call? Generally the only reason you use low precision KV is to reduce HBM util and as such you would precompute your scales during KV quantization.

cthi · 2026-08-11T13:25:47Z

+
+    if use_fp8_kv:
+        # Opt-in fp8-KV: quantize dense KV to fp8 per call (lossy, gfx950 only).
+        from .flydsl.pa_decode_fp8_dispatch import is_fp8_paged_decode_available


we prefer not to use method level import if possible

cthi · 2026-08-11T13:26:01Z

+    from .flydsl.layout_utils import canonicalize_qkv_5d, normalize_seq_positions
+    from .flydsl.pa_decode_dense import pa_decode_launch


we prefer not to use method level import if possible

cthi · 2026-08-11T13:26:58Z

    SUPPORTS_DROPOUT = False
    SUPPORTS_CUSTOM_SCALE = True
    SUPPORTS_BMGHK = True
-    NAME = "ck_splitKF"


Is it possible we can add FlyDSL as a new backend to fmha first, before removing CK? That will make it simpler to accept the flyDSL portion and test it independently. cc @bottler

cthi · 2026-08-11T13:27:58Z

    )

+    # fp8 (~2 mantissa bits): loosened to absorb single-element grid rounding noise
+    # (gfx950 OCP e4m3fn snaps a few values differently than the fnuz grid).


Can you loosen this for ROCm only?

cthi · 2026-08-11T13:28:49Z

        configs=TRITON_CONFIGS,
        key=AUTOTUNER_KEY,
-        use_cuda_graph=True,
+        use_cuda_graph=False if torch.version.hip else True,


If this is broken on gfx950, we can just gate it as such? e.g. leave the current behavior for gfx942?

cthi · 2026-08-11T13:35:14Z

+    return do_bench(fn, (), BenchOptions(cuda_graph=False, rep_ms=rep_ms))
+
+
+class EmptyGraphError(RuntimeError):


Can all this EmptyGraph logic be removed now? And you can just use standard triton.do_bench_cudagraph ?

cthi · 2026-08-11T13:35:30Z

+        from mslk.attention.fmha.flydsl.pa_decode_dense import pa_decode_launch
+        from mslk.flydsl.common import is_flydsl_available


please prefer top levle imports if possible

cthi · 2026-08-11T13:36:23Z

+    return ms
+
+
+def _bench_ms(fn: Callable, rep_ms: int = 200, use_cuda_graph: bool = False) -> float:


Why cuda graph=True is not default

meta-cla Bot added the cla signed label Jul 29, 2026

avbokovoy marked this pull request as ready for review July 29, 2026 10:43

facebook-github-tools Bot added the module: rocm label Jul 29, 2026

cthi reviewed Aug 7, 2026

View reviewed changes

avbokovoy added 3 commits August 10, 2026 08:27

Fix format

f18612c

Trim commits, fix cuda-graph capture in benchmarks, refactor utils

149c11a

avbokovoy force-pushed the flydsl-decode-fp8-graph branch from 930260d to 149c11a Compare August 10, 2026 12:24

cthi requested changes Aug 11, 2026

View reviewed changes


		"""Shared low-level FlyDSL helpers for attention kernels.

		Only pip `flydsl==0.2.2` is imported (no ~/FlyDSL/kernels imports).


		# pyre-strict

		"""Shared low-level FlyDSL helpers for attention kernels.

		from .flydsl.layout_utils import canonicalize_qkv_5d, normalize_seq_positions
		from .flydsl.pa_decode_dense import pa_decode_launch

		return do_bench(fn, (), BenchOptions(cuda_graph=False, rep_ms=rep_ms))


		class EmptyGraphError(RuntimeError):

		from mslk.attention.fmha.flydsl.pa_decode_dense import pa_decode_launch
		from mslk.flydsl.common import is_flydsl_available

		return ms


		def _bench_ms(fn: Callable, rep_ms: int = 200, use_cuda_graph: bool = False) -> float:

Uh oh!

Conversation

avbokovoy commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Kernels

Performance

Other changes

Test plan

Uh oh!

meta-codesync Bot commented Aug 5, 2026

Uh oh!

cthi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

avbokovoy commented Aug 10, 2026

Uh oh!

cthi left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

avbokovoy commented Jul 29, 2026 •

edited

Loading

cthi left a comment •

edited

Loading